# flake8: noqa: E501
import os
from typing import Optional, Type

from asyncer import asyncify

from lagent.actions.base_action import AsyncActionMixin, BaseAction, tool_api
from lagent.schema import ActionReturn, ActionStatusCode
from .parser import BaseParser, JsonParser


class GoogleScholar(BaseAction):
    """
    谷歌学术搜索插件。
    用于搜索和获取学术论文、作者信息等。
    
    Args:
        api_key: Serper API 密钥，可在 https://serper.dev 免费获取
        description: 工具描述
        parser: 结果解析器类
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        description: Optional[dict] = None,
        parser: Type[BaseParser] = JsonParser,
    ):
        super().__init__(description, parser)
        # 从环境变量或参数获取 API 密钥
        api_key = os.environ.get('SERPER_API_KEY', api_key)
        if api_key is None:
            raise ValueError('请设置 Serper API 密钥(环境变量 SERPER_API_KEY 或参数 api_key)')
        self.api_key = api_key

    @tool_api(explode_return=True)
    def search_google_scholar(
        self,
        query: str,  # 搜索查询
        cites: Optional[str] = None,  # 引用搜索的文章ID
        as_ylo: Optional[int] = None,  # 起始年份
        as_yhi: Optional[int] = None,  # 结束年份
        scisbd: Optional[int] = None,  # 最近文章过滤
        cluster: Optional[str] = None,  # 所有版本搜索的文章ID
        hl: Optional[str] = None,  # 界面语言
        lr: Optional[str] = None,  # 限制语言
        start: Optional[int] = None,  # 结果偏移量
        num: Optional[int] = None,  # 返回结果数量
        as_sdt: Optional[str] = None,  # 搜索类型/过滤器
        safe: Optional[str] = None,  # 成人内容过滤级别
        filter: Optional[str] = None,  # 相似/省略结果过滤
        as_vis: Optional[str] = None,  # 是否包含引用
    ) -> dict:
        """
        基于查询搜索学术文章。

        Returns:
            dict: 文章信息
                - title: 三篇选定论文的标题列表
                - cited_by: 三篇论文的引用数列表
                - organic_id: 三篇论文的结果ID列表
                - pub_info: 发表信息
        """
        from serpapi import GoogleSearch

        # 构建搜索参数
        params = {
            'q': query,
            'engine': 'google_scholar',
            'api_key': self.api_key,
            'cites': cites,
            'as_ylo': as_ylo,
            'as_yhi': as_yhi,
            'scisbd': scisbd,
            'cluster': cluster,
            'hl': hl,
            'lr': lr,
            'start': start,
            'num': num,
            'as_sdt': as_sdt,
            'safe': safe,
            'filter': filter,
            'as_vis': as_vis,
        }

        try:
            # 执行搜索
            search = GoogleSearch(params)
            r = search.get_dict()
            results = r['organic_results']
            
            # 提取结果信息
            title = []
            snippets = []
            cited_by = []
            organic_id = []
            pub_info = []
            
            # 处理前三个结果
            for item in results[:3]:
                title.append(item['title'])
                pub_info.append(item['publication_info']['summary'])
                citation = item['inline_links'].get('cited_by', {'total': ''})
                cited_by.append(citation['total'])
                snippets.append(item['snippet'])
                organic_id.append(item['result_id'])
                
            return dict(
                title=title,
                cited_by=cited_by,
                organic_id=organic_id,
                snippets=snippets
            )
            
        except Exception as e:
            return ActionReturn(
                errmsg=str(e),
                state=ActionStatusCode.HTTP_ERROR
            )

    @tool_api(explode_return=True)
    def get_author_information(self, author_id: str, **kwargs) -> dict:
        """
        通过作者ID获取作者信息。

        Args:
            author_id: 作者ID
            **kwargs: 其他参数(语言、排序等)

        Returns:
            dict: 作者信息
                * name: 作者姓名
                * affliation: 所属机构
                * articles: 最多3篇文章
                * website: 个人主页URL
        """
        from serpapi import GoogleSearch

        params = {
            'engine': 'google_scholar_author',
            'author_id': author_id,
            'api_key': self.api_key,
            'hl': kwargs.get('hl'),
            'view_op': kwargs.get('view_op'),
            'sort': kwargs.get('sort'),
            'citation_id': kwargs.get('citation_id'),
            'start': kwargs.get('start'),
            'num': kwargs.get('num'),
            'no_cache': kwargs.get('no_cache'),
            'async': kwargs.get('async_req'),
            'output': kwargs.get('output'),
        }
        try:
            search = GoogleSearch(params)
            results = search.get_dict()
            author = results['author']
            articles = results.get('articles', [])
            return dict(
                name=author['name'],
                affiliations=author.get('affiliations', ''),
                website=author.get('website', ''),
                articles=[dict(title=article['title'], authors=article['authors']) for article in articles[:3]],
            )
        except Exception as e:
            return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)

    @tool_api(explode_return=True)
    def get_citation_format(self, q: str, **kwargs) -> dict:
        """
        获取文章的 MLA 引用格式。

        Args:
            q: 文章的结果ID
            **kwargs: 其他参数(缓存、输出格式等)

        Returns:
            dict: 引用格式信息
                * authors: 文章作者
                * citation: 引用格式
        """
        from serpapi import GoogleSearch

        params = {
            'q': q,
            'engine': 'google_scholar_cite',
            'api_key': self.api_key,
            'no_cache': kwargs.get('no_cache'),
            'async': kwargs.get('async_'),
            'output': kwargs.get('output'),
        }
        try:
            search = GoogleSearch(params)
            results = search.get_dict()
            citation = results['citations']
            citation_info = citation[0]['snippet']
            return citation_info
        except Exception as e:
            return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)

    @tool_api(explode_return=True)
    def get_author_id(self, mauthors: str, **kwargs) -> dict:
        """
        通过作者姓名获取作者ID。

        Args:
            mauthors: 作者姓名
            **kwargs: 其他参数(语言、分页等)

        Returns:
            dict: 作者ID信息
                * author_id: 作者ID
        """
        from serpapi import GoogleSearch

        params = {
            'mauthors': mauthors,
            'engine': 'google_scholar_profiles',
            'api_key': self.api_key,
            'hl': kwargs.get('hl'),
            'after_author': kwargs.get('after_author'),
            'before_author': kwargs.get('before_author'),
            'no_cache': kwargs.get('no_cache'),
            'async': kwargs.get('async_'),
            'output': kwargs.get('output'),
        }
        try:
            search = GoogleSearch(params)
            results = search.get_dict()
            profile = results['profiles']
            author_info = dict(author_id=profile[0]['author_id'])
            return author_info
        except Exception as e:
            return ActionReturn(errmsg=str(e), state=ActionStatusCode.HTTP_ERROR)


class AsyncGoogleScholar(AsyncActionMixin, GoogleScholar):
    """异步版本的谷歌学术搜索插件"""
    
    # 异步版本的方法实现与同步版本类似，
    # 主要是使用 @asyncify 装饰器将同步方法转换为异步方法
    # 这里省略了重复的注释...
